package com.aphrodite.security;

/**
 * An Tiny Encryption Algorithm crypto.
 * <p>
 * Simple and efficient for mobile devices
 * </p>
 */
public class TEACrypto {
	private static final int rounds = 32, BLOCK_SIZE = 8, key_size = 16, delta = 0x9E3779B9, d_sum = 0xC6EF3720; // sum
	// on
	// decrypt
	/*
	 * the expanded key array of 4 subkeys
	 */
	private int _a, _b, _c, _d;
	private boolean _forEncryption;
	private boolean _initialised;

	/**
	 * Create an instance of the TEA encryption algorithm and set some defaults
	 */
	public TEACrypto() {
		_initialised = false;
	}

	private int bytesToInt(byte[] in, int inOff) {
		return ((in[inOff++]) << 24) | ((in[inOff++] & 255) << 16) | ((in[inOff++] & 255) << 8) | ((in[inOff] & 255));
	}

	public byte[] decrypt(final byte[] cipher) {
		return decryptBytes(cipher);
	}

	private int decryptBlock(byte[] in, int inOff, byte[] out, int outOff) {
		// Pack bytes into integers
		int v0 = bytesToInt(in, inOff);
		int v1 = bytesToInt(in, inOff + 4);

		int sum = d_sum;

		for (int i = 0; i != rounds; i++) {
			v1 -= ((v0 << 4) + _c) ^ (v0 + sum) ^ ((v0 >>> 5) + _d);
			v0 -= ((v1 << 4) + _a) ^ (v1 + sum) ^ ((v1 >>> 5) + _b);
			sum -= delta;
		}

		unpackInt(v0, out, outOff);
		unpackInt(v1, out, outOff + 4);

		return BLOCK_SIZE;
	}

	public byte[] decryptBytes(final byte[] source) {
		_forEncryption = false;
		final int length = ((source[2] & 0xFF) * 256) + (source[1] & 0xFF);
		final int blocks = (length / BLOCK_SIZE) + (length % BLOCK_SIZE != 0 ? 1 : 0);

		final byte[] output = new byte[blocks * BLOCK_SIZE];
		int inputOffset = 3;
		int outputOffset = 0;
		for (int i = 0; i < blocks; i++) {
			processBlock(source, inputOffset, output, outputOffset);
			inputOffset += BLOCK_SIZE;
			outputOffset += BLOCK_SIZE;
		}
		final byte[] retval = new byte[length];
		System.arraycopy(output, 0, retval, 0, length);
		return retval;
	}

	public byte[] encrypt(final byte[] source) {
		return encrypt(source, 0, source.length);
	}

	public byte[] encrypt(final byte[] source, final int offset, final int length) {
		_forEncryption = true;

		final int blocks = (length / BLOCK_SIZE) + (length % BLOCK_SIZE != 0 ? 1 : 0);
		final byte[] output = new byte[3 + (blocks * BLOCK_SIZE)];
		final byte[] input = new byte[blocks * BLOCK_SIZE];
		System.arraycopy(source, offset, input, 0, length);
		for (int i = length; i < input.length; i++) {
			input[i] = 0;
		}
		int inputOffset = 0;
		int outputOffset = 3;
		output[0] = 1;
		output[1] = (byte) (length & 0x00FF);
		output[2] = (byte) (length >>> 8);
		for (int i = 0; i < blocks; i++) {
			processBlock(input, inputOffset, output, outputOffset);
			inputOffset += BLOCK_SIZE;
			outputOffset += BLOCK_SIZE;
		}
		return output;
	}

	private int encryptBlock(byte[] in, int inOff, byte[] out, int outOff) {
		// Pack bytes into integers
		int v0 = bytesToInt(in, inOff);
		int v1 = bytesToInt(in, inOff + 4);

		int sum = 0;

		for (int i = 0; i != rounds; i++) {
			sum += delta;
			v0 += ((v1 << 4) + _a) ^ (v1 + sum) ^ ((v1 >>> 5) + _b);
			v1 += ((v0 << 4) + _c) ^ (v0 + sum) ^ ((v0 >>> 5) + _d);
		}

		unpackInt(v0, out, outOff);
		unpackInt(v1, out, outOff + 4);

		return BLOCK_SIZE;
	}

	public String getAlgorithmName() {
		return "TEA";
	}

	public int getBlockSize() {
		return BLOCK_SIZE;
	}

	/**
	 * initialise
	 * 
	 * @param forEncryption
	 *            whether or not we are for encryption.
	 * @param params
	 *            the parameters required to set up the cipher.
	 * @exception IllegalArgumentException
	 *                if the params argument is inappropriate.
	 */
	public void init(boolean forEncryption, byte[] key) {

		_forEncryption = forEncryption;
		_initialised = true;
		setKey(key);
	}

	/**
	 * 
	 * @param in
	 * @param inOff
	 * @param out
	 * @param outOff
	 * @return
	 */
	public int processBlock(byte[] in, int inOff, byte[] out, int outOff) {
		if (!_initialised) {
			throw new IllegalStateException(getAlgorithmName() + " not initialised");
		}

		if ((inOff + BLOCK_SIZE) > in.length) {
			throw new IllegalArgumentException("input buffer too short");
		}

		if ((outOff + BLOCK_SIZE) > out.length) {
			throw new IllegalArgumentException("output buffer too short");
		}

		return (_forEncryption) ? encryptBlock(in, inOff, out, outOff) : decryptBlock(in, inOff, out, outOff);
	}

	public void reset() {}

	/**
	 * Re-key the cipher.
	 * <p>
	 * 
	 * @param key
	 *            the key to be used
	 */
	private void setKey(byte[] key) {
		_a = bytesToInt(key, 0);
		_b = bytesToInt(key, 4);
		_c = bytesToInt(key, 8);
		_d = bytesToInt(key, 12);
	}

	private void unpackInt(int v, byte[] out, int outOff) {
		out[outOff++] = (byte) (v >>> 24);
		out[outOff++] = (byte) (v >>> 16);
		out[outOff++] = (byte) (v >>> 8);
		out[outOff] = (byte) v;
	}
}
